from volcenginesdkarkruntime import Ark

client = Ark(
    api_key="504bfd41-c285-4924-9107-74bb71846399",
    base_url="https://ark.cn-beijing.volces.com/api/v3",
)

# Non-streaming:
# print("----- standard request -----")
# completion = client.chat.completions.create(
#     model="doubao-1-5-pro-256k-250115",
#     messages = [
#         {"role": "system", "content": "你是人工智能助手"},
#         {"role": "user", "content": "常见的十字花科植物有哪些？"},
#     ],
# )
# print(completion.choices[0].message.content)

# Streaming:
print("----- streaming request -----")
stream = client.chat.completions.create(
    model="doubao-1-5-pro-256k-250115",
    messages = [
        {"role": "system", "content": "你是懂电力经济数的人工智能助手"},
        {"role": "user", "content": "如何利用电力数据预测经济？"},
    ],
    stream=True
)
for chunk in stream:
    if not chunk.choices:
        continue
    print(chunk.choices[0].delta.content, end="")
print()